#ifndef _DssFragmentation_H_
#define _DssFragmentation_H_

//
// DssFragmentation.h
//
// Created by TB on 11/18/11.
// Copyright 2011 BNL. All rights reserved.
//

#include <map>
#include <vector>

#include "tppmc/Enumerations.h"
#include "tppmc/FragmentationFunction.h"
#include "tppmc/KeyedLists.h"
#include "tppmc/Particle.h"
#include "tppmc/TF1J.h"

namespace tppmc {
   
   class DssFragmentation : public FragmentationFunction {
      
   public:
      
      /**
       Constructor.
       Initialise with the order of calculation to use (LO or NLO, defined
       int Enumerations.h) and the mean squared pT from fragmentation.
       */
      DssFragmentation(int order = LO, double meanPt2 = 0.2);
      
      /**
       Destructor.
       */
      virtual ~DssFragmentation();
      
      /**
       Fragment the input parton into a hadron at the requested scale Q2.
       If scale is < 0, the scale is the parton pT.
       Returns a vector with a single Hadron.
       Sets the type, precursor parton and z of the returned hadron.
       All other fields have their default value.
       */
      virtual std::vector<tppmc::Hadron> Fragment(const pythpart&,
                                                  double scale = -1.);

      /**
       Return the value of the fragmentation function for the input parton type
       to the input hadron type, at the z of the hadron and the requested scale
      */
      virtual double Evaluate(const Parton&, const Hadron&, double scale);
      
      /**
       Returns true if the fragmentation of the input parton type
       is supported by this function.
       */
      virtual bool Supports(const pythpart&) const;
      
      virtual std::string Name() const { return "DSS"; }
      virtual double ZMin() const { return 0.05; }
      virtual double ZMax() const { return 1.; }

      virtual double MinimumScale() const { return 2.; }
      
      virtual double MaximumScale() const { return 1.e5; }

protected:
      
      /**
       Enumeration for function parameters.
       */
      enum {
         FUNCTION, SCALE, PARTON, ORDER, HADRON, CHARGE
      };
      
      virtual TF1* Function();
      
      /**
       Set the hadron-type-dependent parameters of the ROOT function
       so that the function is called for the requested hadron type.
       */
      void SetFunctionParameters(int htype);
      
      /**
       Populate the probability tables for each hadron,
       parton and Q2 bin.
       */
      void PopulateKeyedLists();
      
      /**
       Set "error" values for all hadrons in the probability
       table for this parton and Q2.
       */
      void SetErrorValues(int parton, int q2bin);
      
      /**
       Lists hadron species supported by the fragmentation function.
       */
      virtual const std::vector<int>& Hadrons() const;
      
      /**
       Lists parton species supported by the fragmentation function.
       */
      virtual const std::vector<int>& Partons() const;
      
      /**
       Returns the table corresponding to the requested
       hadron type, NULL if there is no table.
       */
      KeyedLists* HadronTable(int);
      
      /**
       Randomly throw a hadron type based on the input parton type and Q2.
       Returns the PDG code of the chosen species.
       */
      virtual int ThrowHadronType(const pythpart&, double Q2);
      
      /**
       Used in initialisation.
       Computes probabilities for fragmenting to each hadron type
       in a single Q2 bin for a single parton type.
       */
      void ComputeHadronProbabilities(int parton, int Q2Bin);
      
      /*
       Add a KeyedLists, properly populated for the
       number of Q2 bins and parton types supported,
       for a specified hadron species.
       In each KeyedList:
         The key value is a parton type
         The list represents Q2 bins, each storeing the probability for this
         parton in this Q2 bin to fragment into this hadron type.
       */
      void Insert(int hadron);

      typedef std::map<int, KeyedLists> Table;
      
      // Each key is a hadron type.
      // Each value is a keyed list of proabilities for fragmentation
      // into that hadron - keyed by parton, with each list entry
      // representing a different Q2 bin.
      // i.e. essentially a 3D array of [hadron][parton][Q2]
      // but keyed by hadron & parton type.
      Table mTable; //!
      
      int mOrder;
      double mMeanPt2;
      
      static int smNQBins;
      static double smQMin;
      static double smQMax;
      
   private:
   
      // Get a weird crash if this is a TF1J instead of TF1J*...
      TF1J* mFunction; //!
      
      ClassDef(tppmc::DssFragmentation, 1)
   };
   
} // namespace tppmc

#endif
